import logging
import enum
from datetime import datetime
from itertools import groupby
from django.conf import settings
from storage.rawdata import RawEventDataStorage
from common.es_core.es_interface import es_bulk_insert_interface, es_insert_interface

_LOGGER = logging.getLogger(__name__)


class IntervalType(enum.Enum):
    DAY = 0,
    MONTH = 1


class ESBypassDataStorage(RawEventDataStorage):
    """
    ElasticSearch 存储器，将数据按分流类型索引存储到es
    """
    def __init__(self, name: str, cluster='maestro', interval_type=IntervalType.DAY):
        """
        :param name: ES index prefix
        """
        super().__init__(name)
        self.index = name
        assert cluster in settings.ES_CLUSTER_CLIENT_URLS
        self.cluster = cluster
        assert isinstance(interval_type, IntervalType)
        self.interval_type = interval_type
        self.batch_buffer = list()

    def save(self, event):
        es_insert_interface(self._get_es_index(event), event, cluster=self.cluster)

    def add_batch(self, events):
        assert isinstance(events, list) or isinstance(events, dict)
        if isinstance(events, list):
            self.batch_buffer.extend(events)
        else:
            self.batch_buffer.append(events)

    def execute_batch(self):
        if not self.batch_buffer:
            return
        for interval_value, events in groupby(self.batch_buffer, key=self._extract_interval_str_from_event):
            self._bulk_save(interval_value, list(events))
        self.batch_buffer.clear()

    def _bulk_save(self, interval_value: str, events):
        index = '_'.join([self.index, 'log', interval_value])
        es_bulk_insert_interface(index, events, cluster=self.cluster)

    def _get_es_index(self, event):
        assert '_event_time' in event
        return '_'.join([self.index, 'log', self._extract_interval_str_from_event(event)])

    def _extract_interval_str_from_event(self, event):
        assert '_event_time' in event
        event_time = event['_event_time']
        if self.interval_type == IntervalType.DAY:
            return datetime.fromtimestamp(event_time).strftime('%Y_%m_%d')
        else:
            return datetime.fromtimestamp(event_time).strftime('%Y_%m')
